package org.hepeng.workx.spring.cloud.netflix.ribbon.loadbalancer;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;
import com.netflix.loadbalancer.ZoneAvoidanceRule;
import org.apache.commons.collections.CollectionUtils;
import org.hepeng.workx.spring.cloud.netflix.ribbon.RibbonCommandContextHolder;
import org.hepeng.workx.spring.cloud.netflix.ribbon.SafeRibbonCommandContext;

import java.util.List;
import java.util.Objects;

/**
 * @author he peng
 */
public abstract class AbstractRequestConditionalRule extends AbstractLoadBalancerRule {

    protected IRule rule;

    public AbstractRequestConditionalRule() {
        this.rule = new ZoneAvoidanceRule();
    }

    @Override
    public void setLoadBalancer(ILoadBalancer lb) {
        super.setLoadBalancer(lb);
        this.rule.setLoadBalancer(lb);
    }

    @Override
    public void initWithNiwsConfig(IClientConfig clientConfig) {

    }

    @Override
    public Server choose(Object key) {
        SafeRibbonCommandContext ribbonContext = RibbonCommandContextHolder.get();
        if (Objects.isNull(ribbonContext)) {
            return this.rule.choose(key);
        }

        ILoadBalancer lb = getLoadBalancer();
        List<Server> servers = lb.getAllServers();
        if (CollectionUtils.isEmpty(servers)) {
            return null;
        }

        Server server = null;
        if (isSupported(key , ribbonContext)) {
            server = choose(key, ribbonContext, lb);
        }

        return Objects.nonNull(server) ? server : this.rule.choose(key);
    }

    public abstract Server choose(Object key , SafeRibbonCommandContext ribbonContext , ILoadBalancer lb);

    protected abstract boolean isSupported(Object key , SafeRibbonCommandContext ribbonContext);
}
